package com.zaoshui.hudada.manager;

import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.*;
import io.reactivex.Flowable;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;

@Component
public class AIManager {
    @Resource
    private ClientV4 clientV4;

    // 较稳定的随机数
    private static final float STABLE_TEMPERATURE = 0.05f;

    // 不稳定的随机数
    private static final float UNSTABLE_TEMPERATURE = 0.99f;


    /*
    * 同步方法(较稳定)
    * */
    public String doSyncStableRequest(String systemMessage,String userMessage){
        return doRequest(systemMessage,userMessage,Boolean.FALSE,STABLE_TEMPERATURE);
    }

    /*
    * 同步方法(不稳定)
    * */
    public String doSyncUnStableRequest(String systemMessage,String userMessage){
        return doRequest(systemMessage,userMessage,Boolean.FALSE,UNSTABLE_TEMPERATURE);
    }

    /*
    * 同步方法
    * */
    public String doSyncRequest(String systemMessage,String userMessage,Float temperature){
        return doRequest(systemMessage,userMessage,Boolean.FALSE,temperature);
    }

    /*
    * 简化通用请求
    * */
    public String doRequest(String systemMessage,String userMessage,Boolean stream,Float temperature){
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage userChatMessage = new ChatMessage(ChatMessageRole.USER.value(),userMessage);
        ChatMessage systemChatMessage = new ChatMessage(ChatMessageRole.SYSTEM.value(), systemMessage);
        messages.add(systemChatMessage);
        messages.add(userChatMessage);
        return doRequest(messages,stream,temperature);
    }

    /*
    * 通用请求
    * */
    public String doRequest(List<ChatMessage> messages, Boolean stream, Float temperature) {
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(Constants.ModelChatGLM4)
                .stream(Boolean.FALSE)
                .invokeMethod(Constants.invokeMethod)
                .maxTokens(2048)
                .messages(messages)
                .build();
        ModelApiResponse invokeModelApiResp = clientV4.invokeModelApi(chatCompletionRequest);
        return invokeModelApiResp.getData().getChoices().get(0).getMessage().getContent().toString();
    }

    /*
     * 简化通用流式请求
     * */
    public Flowable<ModelData> doStreamRequest(String systemMessage, String userMessage, Float temperature){
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage userChatMessage = new ChatMessage(ChatMessageRole.USER.value(),userMessage);
        ChatMessage systemChatMessage = new ChatMessage(ChatMessageRole.SYSTEM.value(), systemMessage);
        messages.add(systemChatMessage);
        messages.add(userChatMessage);
        return doStreamRequest(messages,temperature);
    }

    /*
     * 通用流式请求
     * */
    public Flowable<ModelData> doStreamRequest(List<ChatMessage> messages, Float temperature) {
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(Constants.ModelChatGLM4)
                .stream(Boolean.TRUE)
                .invokeMethod(Constants.invokeMethod)
                .maxTokens(2048)
                .messages(messages)
                .build();
        ModelApiResponse invokeModelApiResp = clientV4.invokeModelApi(chatCompletionRequest);
        return invokeModelApiResp.getFlowable();
    }

}
